# coding: utf-8
# cf.http://d.hatena.ne.jp/white_wheels/20100327/p3
import matplotlib.pylab as plt
import numpy as np


# 梯度
def _numerical_gradient_no_batch(f, x):
    h = 1e-4  # 0.0001
    grad = np.zeros_like(x)  # 创建x同规模的0矩阵

    for idx in range(x.size):
        tmp_val = x[idx]  # 临时用一下数组中的第一个x值
        x[idx] = float(tmp_val) + h
        fxh1 = f(x)  # f(x+h)

        x[idx] = tmp_val - h
        fxh2 = f(x)  # f(x-h)
        grad[idx] = (fxh1 - fxh2) / (2 * h)

        x[idx] = tmp_val  # 还原原数组的x值

    return grad  # 返回grad数组


# 梯度
def numerical_gradient(f, X):
    if X.ndim == 1:  # 输入的X是一维数组向量，数组两个数字分别表示x0和x1的输入值
        return _numerical_gradient_no_batch(f, X)
    else:
        grad = np.zeros_like(X)

        for idx, x in enumerate(X):
            grad[idx] = _numerical_gradient_no_batch(f, x)

        return grad


def function_2(x):
    if x.ndim == 1:
        return np.sum(x ** 2)
    else:
        return np.sum(x ** 2, axis=1)  # 这句话什么意思？

if __name__ == '__main__':
    x0 = np.arange(-2, 2.5, 0.25)
    x1 = np.arange(-2, 2.5, 0.25)
    X, Y = np.meshgrid(x0, x1)  # 使用meshgrid产生格点矩阵

    X = X.flatten()  # 将数据一维化处理
    Y = Y.flatten()

    grad = numerical_gradient(function_2, np.array([X, Y]))

    plt.figure()

    # 剑鞘
    plt.quiver(X, Y, -grad[0], -grad[1], angles="xy", color="#666666")  # ,headwidth=10,scale=40,color="#444444")

    plt.xlim([-2, 2])
    plt.ylim([-2, 2])
    plt.xlabel('x0')
    plt.ylabel('x1')

    plt.grid()  # 将网络绘画出来
    plt.legend()  # 图例展示
    plt.draw()
    plt.show()
